package BackTrack.Hard;

import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;

public class LC0051 {
    List<List<String>> res = new ArrayList<>();
    boolean[] selected;
    int n;

    public List<List<String>> solveNQueens(int n) {
        this.n = n;
        selected = new boolean[n];
        backTrackEarlyCheck(null);
        return res;
    }

    /**
     * 有限剪枝：在决策树的结点上向下扩展时，进行横向和纵向检查，并和上一个皇后进行横向检查。
     * 耗时17ms。
     */
    private void backTrack(LinkedList<Integer> path) {
        boolean expandFlag = false;
        if (path == null) path = new LinkedList<>();
        int lastSelectedColumn = -1;
        if (!path.isEmpty()) lastSelectedColumn = path.getLast();
        for (int i = 0; i < n; i++) {
            if (selected[i]) continue;
            if (i + 1 < n && i + 1 == lastSelectedColumn) continue;
            if (i - 1 >= 0 && i - 1 == lastSelectedColumn) continue;

            expandFlag = true;
            path.add(i);
            selected[i] = true;
            backTrack(path);
            selected[i] = false;
            path.removeLast();
        }

        if (!expandFlag && checkValidity(path)) {
            res.add(generateReturnValue(path));
        }
    }

    /**
     * 在决策树的结点上向下扩展时检查合法性
     * 该方法耗时18ms。
     */
    private void backTrackEarlyCheck(LinkedList<Integer> path) {
        boolean expandFlag = false;
        if (path == null) path = new LinkedList<>();
        int lastSelectedColumn = -1;
        if (!path.isEmpty()) lastSelectedColumn = path.getLast();
        for (int i = 0; i < n; i++) {
            if (selected[i]) continue;
            int row = path.size(); // 当前的行号
            // 斜向检查
            boolean obliqueValidity = true;
            for (int j = 0; j < row; j++) {
                int delta = row - j;
                if (path.get(j) == i - delta || path.get(j) == i + delta) {
                    obliqueValidity = false;
                    break;
                }
            }
            if (!obliqueValidity) continue;

            expandFlag = true;
            path.add(i);
            selected[i] = true;
            backTrack(path);
            selected[i] = false;
            path.removeLast();
        }

        if (!expandFlag) {
            res.add(generateReturnValue(path));
        }
    }

    private boolean checkValidity(List<Integer> path) {
        if (path.size() < n) return false;

        for (int row = 0; row < n; row++) {
            int column = path.get(row);
            // 不必进行横向纵向检查，生成路径时的算法可以确保不会在横向和纵向上冲突
            // 斜向检查
            for (int i = 1; i < n; i++) {
                if (row - i >= 0 && column - i >= 0) {
                    if (path.get(row - i) == (column - i)) return false;
                }
                if (row + i < n && column - i >= 0) {
                    if (path.get(row + i) == (column - i)) return false;
                }
                if (row - i >= 0 && column + i < n) {
                    if (path.get(row - i) == (column + i)) return false;
                }
                if (row + i < n && column + i < n) {
                    if (path.get(row + i) == (column + i)) return false;
                }
            }
        }

        return true;
    }

    private List<String> generateReturnValue(List<Integer> pickedColumns) {
        List<String> ret = new ArrayList<>();
        for (int column : pickedColumns) {
            char[] charArray = new char[n];
            for (int i = 0; i < n; i++) {
                if (i == column) charArray[i] = 'Q';
                else charArray[i] = '.';
            }
            ret.add(new String(charArray));
        }
        return ret;
    }

    public static void main(String[] args) {
        LC0051 solution = new LC0051();
        solution.solveNQueens(4);
    }
}
